import os

import PIL
import numpy as np
# from PIL.Image import Image

os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'  # 只显示 warning 和 Error

from tensorflow.python.keras.utils import Sequence

import string

characters = string.digits + string.ascii_uppercase
# print(characters)




class BaseSequence(Sequence):
    """
    基础的数据流生成器，每次迭代返回一个batch
    BaseSequence可直接用于fit_generator的generator参数
    fit_generator会将BaseSequence再次封装为一个多进程的数据流生成器
    而且能保证在多进程下的一个epoch中不会重复取相同的样本
    """

    def __init__(self, paths, characters, batch_size, steps, n_len=4, width=90, height=30):
        self.characters = characters
        self.batch_size = batch_size
        self.steps = steps
        self.n_len = n_len
        self.width = width
        self.height = height
        self.n_class = len(characters)
        self.path = paths
        self.foo = self.gen_batch(paths, batch_size)  # 初始化生成方法

    def __len__(self):
        return self.steps

    # def gen_batch(self):
    #     return

    @staticmethod
    def gen_batch(paths, batch_size=10):
        '''静态方法，用来遍历文件夹'''
        print("starting...iter...")
        print(paths)
        count = 0
        inc = 0
        imgs = []
        lbs = []
        while True:  # 死循环生成
            for path in paths:
                for file in os.listdir(path):
                    file_path = os.path.join(path, file)
                    # if os.path.isdir(file_path):
                    if os.path.isfile(file_path) and file_path.endswith(".jpg"):
                        count += 1
                        inc += 1
                        imgs.append(file_path)
                        lbs.append(file.split("_")[1].split(".")[0].upper())
                        if inc == batch_size:
                            yield imgs, lbs  # 返回当前生成的这批号码
                            inc = 0
                            imgs = []
                            lbs = []
                        # print("res:", res)

    def __getitem__(self, idx):
        # return next(self.foo) #返回下一批,文件名和标签
        imgs, lbs = next(self.foo)  # 返回下一批,文件名和标签
        X = np.zeros((self.batch_size, self.height, self.width, 3), dtype=np.float32)
        y = [np.zeros((self.batch_size, self.n_class), dtype=np.uint8) for i in range(self.n_len)]
        for i in range(self.batch_size):
            random_str = lbs[i]
            image_open = PIL.Image.open(imgs[i])
            X[i] = np.array(image_open) / 255.0 #读取图片文件名，并转换成np array,再除255.0
            for j, ch in enumerate(random_str):
                y[j][i, :] = 0
                y[j][i, self.characters.find(ch)] = 1
        return X, y
    # 重写的父类Sequence中的on_epoch_end方法，在每次迭代完后调用。
    # def on_epoch_end(self):
    #     每次迭代后重新打乱训练集数据
    # np.random.shuffle(self.x_y)

if __name__=='__main__':
    width, height, n_len, n_class = 90, 30, 4, len(characters)
    img_paths = [r'./img_train', r'./img_test']
    img_paths_train = [r'E:\szjj_gzh.261.szjjgateway_01.train']
    img_paths_test = [r'E:\szjj_gzh.261.szjjgateway_01.test']
    animals = BaseSequence(paths=img_paths_train, characters=characters, batch_size=2, steps=2, n_len=n_len, width=width, height=height)
    animals2 = BaseSequence(paths=img_paths_test, characters=characters, batch_size=2, steps=2, n_len=n_len, width=width, height=height)
    for i in range(10):
        print(animals.__getitem__(i))
        print(animals2.__getitem__(i))
